#include <Eigen/Dense>
#include <Eigen/IterativeLinearSolvers>
#include <iostream>
#define MAX_ITER 400
#define atol 1e-8

using namespace std;

int main()
{
    // 实现共轭梯度法算例
    int n = 5; // A: n*n

    Eigen::MatrixXd A(5,5);
    A << 5, 1, 0, 0, 0,
        1, 4, 1, 0, 0,
        0, 1, 3, 1, 0,
        0, 0, 1, 2, 1,
        0, 0, 0, 1, 1;
    cout << "A = \n"
         << A << endl;
    Eigen::VectorXd b(5);
    b << 7, 12, 15, 16, 9;

    // eigen库 cg算法结果
    // Eigen::ConjugateGradient<Eigen::MatrixXd> cg;
    // cg.compute(A);
    // Eigen::VectorXd x_eigen = cg.solve(b);
    // cout << "Eigen Result by CG = \n" << x_eigen << endl;

    // 初始化
    Eigen::VectorXd x = Eigen::VectorXd::Zero(n);
    Eigen::VectorXd r0 = b - A*x, r1;
    
    Eigen::VectorXd d = r0;
    double res = 1e30, alpha=0;

    // 迭代
    for (int i = 0; i < MAX_ITER; i++)
    {
        
        alpha = r0.dot(r0) / d.dot(A * d);
        x += alpha * d;
        // 很神奇! 这里不用auto 就是对的
        r1 = r0 - alpha * A * d;
        res = r1.norm();
        printf("Iter %i, residual = %e.\n", i, res);
        if (res < atol)
            break;
        // 矩阵之间的 / 运算符执行的是 矩阵除法（本质是求解线性方程组 AX = B） 
        auto beta = r1.dot(r1) / r0.dot(r0);
        d = r1 + beta * d;
        r0 = r1;
        
    }
    
    cout << "Numerical Solution = \n"
         << x << endl;
}

